package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"os"
	"proxy/log"
	sn "proxy/utils/signal"
	"sync"
	"time"
)

const PORT = 8000
const (
	binance_spot    = "api.binance.com"
	binance_futures = "fapi.binance.com"
	okx             = "www.okx.com"
	bybit           = "api.bybit.com"
	bitget          = "api.bitget.com"
	htx_spot        = "api.huobi.pro"
	htx_futures     = "api.hbdm.com"
	kucoin_spot     = "api.kucoin.com"
	kucoin_futures  = "api-futures.kucoin.com"
	gateio          = "api.gateio.ws"
)

func main() {
	wait := &sync.WaitGroup{}
	ctx, cancel := context.WithCancel(context.Background())

	p := NewPorxy(wait)
	t := time.Now()

	sn.HandleSignals(func(o os.Signal) {
		cancel()

		wait.Wait()
	})

	log.Error("start on:", PORT, t.Format("2006-01-02 15:04:05"))
	if err := p.ListenAndServe(ctx, PORT); err != nil {
		log.Error(err)
	}
	t = time.Now()
	log.Error("stop on:", t.Format("2006-01-02 15:04:05"))
}

type Proxy struct {
	ctx    context.Context
	Server *http.Server
	wait   *sync.WaitGroup
}

func NewPorxy(wait *sync.WaitGroup) *Proxy {
	p := &Proxy{
		wait: wait,
	}
	p.Server = &http.Server{
		Handler:      p,
		TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)),
	}

	return p
}

func (p *Proxy) ListenAndServe(ctx context.Context, port int) error {
	c := make(chan error)
	go func() {
		p.Server.Addr = fmt.Sprintf(":%d", port)
		c <- p.Server.ListenAndServe()
	}()
	p.ctx = ctx
	select {
	case <-ctx.Done():
		p.Server.Shutdown(context.TODO())
		<-c
		return ctx.Err()
	case err := <-c:
		return err
	}
}

func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) {
	p.wait.Add(1)
	defer p.wait.Done()

	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	//log.Error(r.URL.String())
	client := &http.Client{}
	client.Transport = tr
	defer client.CloseIdleConnections()

	req, err := http.NewRequest(r.Method, r.URL.RequestURI(), r.Body)
	req.URL.Host = r.URL.Host
	req.URL.Scheme = r.URL.Scheme
	req.Method = r.Method

	CopyHeader(req.Header, r.Header)
	if err != nil {
		log.Error(err)
	}
	resp, err := client.Do(req)

	if err != nil {
		log.Error(err)
		w.WriteHeader(404)
		w.Write([]byte("Proxy error: " + err.Error()))
		return
	}

	defer resp.Body.Close()

	CopyHeader(w.Header(), resp.Header)

	w.WriteHeader(resp.StatusCode)
	if _, err := io.Copy(w, resp.Body); err != nil {
		log.Error(err)
	}
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer func() {
		if e := recover(); e != nil {
			log.Error(e)
		}
	}()

	isErrExchange := true
	exchange := r.Header.Get("x-exchange")
	r.URL.Scheme = "https"
	switch exchange {
	case "bybit":
		r.URL.Host = bybit
		isErrExchange = false
	case "okx":
		r.URL.Host = okx
		isErrExchange = false
	case "binance_spot":
		r.URL.Host = binance_spot
		isErrExchange = false
	case "binance_futures":
		r.URL.Host = binance_futures
		isErrExchange = false
	case "htx_spot":
		r.URL.Host = htx_spot
		isErrExchange = false
	case "htx_futures":
		r.URL.Host = htx_futures
		isErrExchange = false
	case "gateio":
		r.URL.Host = gateio
		isErrExchange = false
	case "bitget":
		r.URL.Host = bitget
		isErrExchange = false
	case "kucoin_spot":
		r.URL.Host = kucoin_spot
		isErrExchange = false
	case "kucoin_futures":
		r.URL.Host = kucoin_futures
		isErrExchange = false
	default:
		log.Error("错误的交易所: %s 地址：%+v", exchange, r.URL)
	}

	if !isErrExchange {
		r.Header.Del("x-exchange")
		p.handleHTTP(w, r)
	}

}

func CopyHeader(dst, src http.Header) {
	for k, v := range src {
		dst[k] = v
	}
}
